package mosdi.paa;

import java.util.Arrays;

/** Abstract base class for an probabilistic arithmetic automaton. This class 
 *  precomputes and caches all transition probabilities, emission probabilities
 *  and all arithmetic computations. If this behaviour is not desired, the methods
 *  getTargets(), getTargetProbabilities(), getEmissions(), getEmissionProbabilities(),
 *  and performOperationFast() must be overwritten.
 */
public abstract class PAA {
	/** Preprocessed table of target states to be used 
	 *  in getTargets(). */
	private int[][] targets;
	/** Preprocessed table of transition probabilities to be used 
	 *  in getTargetProbabilities(). */
	private double[][] targetProbabilities;
	/** Preprocessed table of emission probabilities to be used 
	 *  in getEmissions(). */
	private int[][] emissions;
	/** Preprocessed table of emission probabilities to be used 
	 *  in getEmissionProbabilities(). */
	private double[][] emissionProbabilities;
	
	/** The size of set Q. */
	public abstract int getStateCount();
	/** The size of set N. */
	public abstract int getValueCount();
	/** The size of set E. */
	public abstract int getEmissionCount();
	
	public abstract int getStartValue();
	public abstract int getStartState();
	
	/** Probability of going from state q to state p. */
	public abstract double transitionProbability(int state, int targetState);

	/** Probability that states q emits value z. */
	public abstract double emissionProbability(int state, int emission);
	
	public abstract int performOperation(int state, int value, int emission);
	
	/** Returns an array of targets of a given state.
	 *  j in getTargets(i) iff transitionProbability(i,j)>0.0.
	 */
	protected int[] getTargets(int state) {
		if (targets==null) {
			targets = new int[getStateCount()][];
			for (int i=0; i<getStateCount(); ++i) {
				targets[i] = new int[getStateCount()];
				int n = 0;
				for (int j=0; j<getStateCount(); ++j) {
					if (transitionProbability(i,j)>0.0) targets[i][n++] = j;
				}
				targets[i] = Arrays.copyOf(targets[i], n);
			}
		}
		return targets[state];
	}
	
	/** Returns array of transition probabilities for outgoing from 
	 *  given state. The returned array is consistent with the result
	 *  of getTargets(), i.e. getTargetProbabilities(q)[i] is the probability
	 *  if going from state q to state getTargets(q)[i].
	 */
	protected double[] getTargetProbabilities(int state) {
		if (targetProbabilities==null) {
			targetProbabilities = new double[getStateCount()][];
			for (int i=0; i<getStateCount(); ++i) {
				int[] t = getTargets(i);
				targetProbabilities[i] = new double[t.length];
				int n = 0;
				for (int j : t) targetProbabilities[i][n++] = transitionProbability(i,j);
			}
		}
		return targetProbabilities[state];
	}
	
	/** Returns an array of possible emissions of a given state, i.e.
	 *  z in getEmissions(q) iff emissionProbability(q,z)>0.0.
	 */
	protected int[] getEmissions(int state) {
		if (emissions==null) {
			emissions = new int[getStateCount()][];
			for (int i=0; i<getStateCount(); ++i) {
				emissions[i] = new int[getStateCount()];
				int n = 0;
				for (int e=0; e<getEmissionCount(); ++e) {
					if (emissionProbability(i,e)>0.0) emissions[i][n++] = e;
				}
				emissions[i] = Arrays.copyOf(emissions[i], n);
			}
		}
		return emissions[state];
	}

	/** Returns array of emission probabilities of given state 
	 *  The returned array is consistent with the result
	 *  of getEmissions(), i.e. getEmissionProbabilities(q)[i] the 
	 *  probability of the emission e = getEmissions(q)[i].
	 */
	protected double[] getEmissionProbabilities(int state) {
		if (emissionProbabilities==null) {
			emissionProbabilities = new double[getStateCount()][];
			for (int i=0; i<getStateCount(); ++i) {
				int[] e = getEmissions(i);
				emissionProbabilities[i] = new double[e.length];
				int n = 0;
				for (int z : e) emissionProbabilities[i][n++] = emissionProbability(i,z);
			}
		}
		return emissionProbabilities[state];
	}
	
	public double[][] stateValueStartDistribution() {
		double[][] result = new double[getStateCount()][getValueCount()];
		result[getStartState()][getStartValue()] = 1.0;
		return result;
	}
	
	public double[][] computeStateValueDistribution(int iterations) {
		double[][] table1 = stateValueStartDistribution();
		// iterate over all times
		for  (int t=0; t<iterations; ++t) {
			double[][] table2 = new double[getStateCount()][getValueCount()];
			// iterate over all source states
			for (int sourceState=0; sourceState<getStateCount(); ++sourceState) {
				int[] targetStates = getTargets(sourceState);
				double[] probs = getTargetProbabilities(sourceState);
				// iterate over all target states that can be reached from source state
				for (int i=0; i<targetStates.length; ++i) {
					int targetState = targetStates[i];
					double transitionProb = probs[i];
					// use simplified recurrence of emissions are deterministic
					if (this instanceof DeterministicEmitter) {
						int emission = ((DeterministicEmitter)this).getEmission(targetState);
						// iterate over all possible former values
						for (int value=0; value<getValueCount(); ++value) {
							table2[targetState][performOperation(targetState, value, emission)]+=
								table1[sourceState][value] * transitionProb;
						}
					} else {
						int[] emissions = getEmissions(targetState);
						double[] emissionProbs = getEmissionProbabilities(targetState);
						// iterate over all emissions that can be emitted by target state
						for (int j=0; j<emissions.length; ++j) {
							int emission = emissions[j];
							double emissionProb = emissionProbs[j];
							// iterate over all possible former values
							for (int value=0; value<getValueCount(); ++value) {
								table2[targetState][performOperation(targetState, value, emission)]+=
									table1[sourceState][value] * transitionProb * emissionProb;
							}
						}
					}
				}
			}
			table1 = table2;
		}
		return table1;
	}
	
	public double[] computeValueDistribution(int iterations) {
		double[][] dist = computeStateValueDistribution(iterations);
		double[] result = new double[getValueCount()];
		for (int value=0; value<getValueCount(); ++value) {
			for (int state=0; state<getStateCount(); ++state) {
				result[value]+=dist[state][value];
			}
		}
		return result;
	}
	
}
